#include <utility>
#include "common/common.hpp"
#include "opencv2/opencv.hpp"
#include "prep_img_ipp.h"
#include "model_infer_vino.h"
#include "common/device_utility.h"
#include <fstream>

// #include <ipp.h>
// #include <ipp/ippcv.h>
// #include <ipp/ipps.h>

// #include <ipp/ippi.h>
#include <vector>
#include <iostream>

namespace vino_v6
{
	/**
	 * @brief 对检测结果进行非极大值抑制（NMS）以去除冗余框
	 *
	 * @param detections 检测结果的向量，每个元素包含检测框及其置信度
	 * @param nms_threshold_ 非极大值抑制的阈值，IOU超过该值的框会被抑制
	 *
	 * 该函数首先按照检测框的置信度从高到低排序，然后依次比较每一个检测框
	 * 和后面的框，计算它们的IOU，如果IOU超过阈值且它们属于相同类别，
	 * 则将后一个框的置信度设置为0，表示抑制该框。最后移除所有置信度为0的框。
	 */
	void NmsDetect(
		vector<Yolo_VINO::DetRes_VINO> & detections,
		const float nms_threshold_)
	{
		// 按检测框的概率从高到低进行排序
		sort(detections.begin(), detections.end(), [=](
			const Yolo_VINO::DetRes_VINO &left,
			const Yolo_VINO::DetRes_VINO &right) {
			return left.prob > right.prob; // 比较器，确保概率高的在前
		});

		// 遍历检测框列表，对每个检测框与其他检测框比较重叠度
		for (int i = 0; i < (int)detections.size(); i++) {
			for (int j = i + 1; j < (int)detections.size(); j++) {
				// 如果两个检测框属于同一类别
				if (detections[i].classes == detections[j].classes) {
					// 计算两个检测框的IoU值
					float iou = IOUCalculate(detections[i], detections[j]);
					// 如果IoU值超过设定的阈值，则认为这两个检测框重叠度过高
					if (iou > nms_threshold_) {
						// 将重叠的检测框的概率设置为0，这样在后续步骤中它会被移除
						detections[j].prob = 0;
					}
				}
			}
		}

		// 移除所有概率为0的检测框，这些检测框是因为重叠度过高而被标记为0的
		detections.erase(std::remove_if(
			detections.begin(),
			detections.end(),
			[](const Yolo_VINO::DetRes_VINO &det) {
			return det.prob == 0; // 移除概率为0的检测框
		}), detections.end());
	}

	/**
	 * @brief 计算两个检测框之间的交并比（IoU）。
	 *
	 * @param det_a 第一个检测框的结果。
	 * @param det_b 第二个检测框的结果。
	 * @return 两个检测框的IoU值，范围在0到1之间，表示重叠程度。
	 */
	float IOUCalculate(
		vector<Yolo_VINO::DetRes_VINO>::const_reference det_a,
		vector<Yolo_VINO::DetRes_VINO>::const_reference det_b)
	{
		// 计算检测框的左上角和右下角坐标
		float left = std::max(det_a.x - det_a.w / 2, det_b.x - det_b.w / 2);
		float top = std::max(det_a.y - det_a.h / 2, det_b.y - det_b.h / 2);
		float right = std::min(det_a.x + det_a.w / 2, det_b.x + det_b.w / 2);
		float bottom = std::min(det_a.y + det_a.h / 2, det_b.y + det_b.h / 2);

		// 计算交集的宽度和高度
		float width = std::max(0.0f, right - left);
		float height = std::max(0.0f, bottom - top);

		// 计算交集面积
		float inter_area = width * height;

		// 计算两个检测框的面积
		float area_a = det_a.w * det_a.h;
		float area_b = det_b.w * det_b.h;

		// 计算并集面积
		float union_area = area_a + area_b - inter_area;

		// 计算IoU
		return inter_area / union_area;
	}

	/**
	 * @brief 从文件中读取数据
	 *
	 * @param file_path 文件的路径。
	 * @return 返回包含文件内容的std::vector<char>。
	 */
	std::vector<char> get_file_data(const std::string& file_path) {
		// 以二进制模式打开文件
		std::ifstream file(file_path, std::ios_base::binary);
		// 确保文件打开成功
		assert(file.good());
		// 移动文件指针到文件末尾，以获取文件长度
		file.seekg(0, std::ios_base::end);
		// 获取文件长度
		const int file_len = file.tellg();
		// 将文件指针移回文件开头
		file.seekg(0, std::ios_base::beg);
		// 创建一个字符向量，大小为文件长度，初始化为0
		std::vector<char> data(file_len, 0);
		// 读取文件内容到向量中
		file.read(data.data(), file_len);
		// 关闭文件
		file.close();
		// 返回包含文件数据的向量
		return data;
	}

	/**
	 * @brief 对数据进行简单的解密操作。
	 * @param data 待解密的数据。
	 */
	void decrypt_data(std::vector<char>& data) {
		char key = 0x1; // 与加密相同的密钥
		for (size_t i = 0; i < data.size(); ++i) {
			data[i] ^= key; // 异或操作进行解密
		}
	}
	/**
	 * @brief 解密文件并提取XML和BIN数据
	 *
	 * @param merged_path 加密文件的路径
	 * @param info 用于存储解密后的加密信息的结构体
	 * @param xml_data 用于存储解密后的XML数据的向量
	 * @param bin_data 用于存储解密后的BIN数据的向量
	 * @return 如果解密和提取成功，返回true；否则返回false
	 *
	 * 该函数首先读取指定路径的加密文件，并尝试解密文件内容。解密后，函数会检查文件
	 * 数据是否包含完整的加密信息结构体。如果文件数据有效，函数会根据加密信息中的
	 * xml_size和bin_size成员，分别提取XML和BIN数据。在vinoIR模型中，XML数据位于BIN
	 * 数据之前，即文件中的顺序是XML在前，BIN在后。
	 */
	bool decryptFileTo_XB_buffer(
		const std::string& merged_path,
		st_encrypt_info_ovino& info,
		std::vector<char>& xml_data,
		std::vector<char>& bin_data)
	{
		// 检查文件是否是 .hdats 文件
		if (merged_path.size() < 6 || merged_path.substr(merged_path.size() - 6) != ".hdats") {
			throw std::runtime_error("The provided file is not an .hdats file.");
		}

		std::vector<char> merged_data = get_file_data(merged_path);

		// 检查合并数据是否至少包含头部信息
		if (merged_data.size() < sizeof(st_encrypt_header)) {
			throw std::runtime_error("Merged data is too small to contain the header.");
		}

		// 读取头部信息
		st_encrypt_header header{};
		std::memcpy(&header, merged_data.data(), sizeof(st_encrypt_header));

		// 检查合并数据是否包含完整的结构体和文件数据
		if (merged_data.size() < sizeof(st_encrypt_header) + header.struct_size + header.xml_size + header.bin_size) {
			throw std::runtime_error("Merged data is too small to contain all parts.");
		}

		// 读取 st_encrypt_info_ovino 结构体
		std::memcpy(&info, merged_data.data() + sizeof(st_encrypt_header), header.struct_size);

		// 提取XML数据并解密
		xml_data.assign(merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size, merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size + header.xml_size);
		decrypt_data(xml_data);

		// 提取BIN数据并解密
		bin_data.assign(merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size + header.xml_size, merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size + header.xml_size + header.bin_size);
		decrypt_data(bin_data);

		return true;
	}


	/**
	 * @brief 将权重数据从 std::vector<char> 转换为 OpenVINO 的 ov::Tensor 对象。
	 *
	 * @param weightsBuffer 包含权重数据的字符向量。
	 * @return 返回一个 ov::Tensor 对象，其数据类型为 u8，大小与权重数据相同。
	 */
	ov::Tensor vector_to_tensor(const std::vector<char>& weightsBuffer) {
		// 确定权重数据的大小
		size_t weightsSize = weightsBuffer.size();
		// 创建一个 u8 类型的 Tensor，其大小与权重数据的大小相同
		ov::Tensor weights(ov::element::u8, {weightsSize});
		// 将权重数据复制到 Tensor 的缓冲区
		std::memcpy(weights.data(), weightsBuffer.data(), weightsSize);
		return weights;
	}

	/**
	 * @brief 使用空数据测试模型推理性能。
	 *
	 * @param infer_request 已创建的推理请求对象。
	 * @param input_shape 输入张量的形状。
	 * @param input_type 输入张量的数据类型。
	 * @return 测试成功返回 true，测试失败返回 false。
	 *
	 * 该函数创建一个与模型输入形状和类型匹配的空输入张量，填充零值，然后执行推理。
	 * 通过测量推理时间来测试模型性能。
	 */
	bool test_inference_with_empty_data(
		ov::InferRequest& infer_request,
		const ov::Shape& input_shape,
		const ov::element::Type& input_type)
	{
		try {
			// 创建一个空的输入张量，与模型输入形状和类型相匹配
			const ov::Tensor input_tensor(input_type, input_shape);
			const std::vector<float> input_data(
				input_tensor.get_shape().size(), 0.0f); // 使用 float 填充零值

			// 将空数据复制到输入张量
			std::memcpy(input_tensor.data<float>(), input_data.data(), input_data.size() * sizeof(float));

			// 将输入张量设置到推理请求
			infer_request.set_input_tensor(input_tensor);

			// 执行推理并测量时间
			const auto start = std::chrono::high_resolution_clock::now();
			infer_request.start_async();
			infer_request.wait_for(std::chrono::milliseconds(1000));
			const auto end = std::chrono::high_resolution_clock::now();

			// 计算推理时间
			std::chrono::duration<double, std::milli> latency = end - start;
			std::cout << "Inference time: " << latency.count() << " ms" << std::endl;

			return true; // 测试成功
		} catch (const std::exception& e) {
			std::cerr << "Inference failed: " << e.what() << std::endl;
			return false; // 测试失败
		}
	}
	// Yolo_VINO类的构造函数，用于初始化对象。
	/**
	 * 初始化Yolo_VINO对象，创建CUDA流以管理异步CUDA操作。
	 */
	Yolo_VINO::Yolo_VINO() = default;

	// Yolo_VINO类的析构函数，用于释放对象占用的资源。
	/**
	 * 释放Yolo_VINO对象占用的所有资源，包括CUDA流、TensorRT执行上下文和引擎。
	 */
	Yolo_VINO::~Yolo_VINO() = default;

	// 重置YoloDet2对象，释放所有已分配的资源。
	/**
	 * 释放YoloDet2对象中所有已分配的资源，包括TensorRT执行上下文、引擎和CUDA设备缓冲区。
	 */
	void Yolo_VINO::Reset() {
		for (auto &req : infer_requests) {
			req = {}; // 将每个推理请求对象设置为默认构造的空对象
		}
		compiled_model = {}; // 将编译后的模型对象设置为默认构造的空对象
		// core = ; // 将OpenVINO核心对象设置为默认构造的空对象
	}

	/**
	 * @brief 初始化Yolo_VINO对象，设置GPU并加载模型。
	 *
	 * @param hdats_path 模型文件的路径。
	 * @return 如果模型加载成功，返回 true；否则返回 false。
	 */
	bool Yolo_VINO::Init(
	    const char* hdats_path)
	{
	    // 步骤 1：初始化 OpenVINO Runtime 核心
	    core = Core(); // 创建 OpenVINO 核心对象

	    // 步骤 2：解密文件&编译模型
	    const std::string file_path = hdats_path; // 将C风格字符串转换为std::string对象。

	    std::cout << "Loading model from " << hdats_path << std::endl;

	    // 分别创建 XML 和 BIN 数据缓冲区
	    std::vector<char> xml_buffer = {};
	    std::vector<char> bin_buffer = {};

	    decryptFileTo_XB_buffer(
	        hdats_path,
	        o_m_info,
	        xml_buffer,
	        bin_buffer);

	    // 假设 grids 是一个存储网格大小的向量
	    std::vector<std::pair<int, int>> grids; // 存储每个步长的网格大小

	    // 根据步长计算网格大小
	    for (int i = 0; i < o_m_info.stride_num; ++i) {
	        // 为每个步长创建一个新的向量，存储该步长的网格尺寸
	        std::vector<int> grid_size(2);
	        // 计算每个步长的网格大小，并存储在grid_size中
	        grid_size[0] = o_m_info.image_height / o_m_info.strides[i]; // 网格高度
	        grid_size[1] = o_m_info.image_width / o_m_info.strides[i]; // 网格宽度
	        // 将grid_size添加到grids向量中
	        m_grids.push_back(grid_size);
	    }

	    // LoadEngine();
	    std::cout << "Loading model ok" << std::endl;

	    std::string model_xml(xml_buffer.begin(), xml_buffer.end());
	    // 将 BIN 数据转换为 ov::Tensor
	    ov::Tensor bin_tensor = vector_to_tensor(bin_buffer);

	    try {
	        // 读取模型文件并创建模型对象
	        std::shared_ptr<ov::Model> model = core.read_model(model_xml,bin_tensor);
	        // 将模型编译到指定设备
	        compiled_model = core.compile_model(model, "AUTO");
	    }
	    catch (const std::exception& e) { // 捕获异常
	        std::cerr << "Model compilation exception: " << e.what() << std::endl; // 输出异常信息
	        throw; // 重新抛出异常
	    }

	    // 步骤 3：创建两个推理请求
	    size_t num_requests = o_m_info.batch_size; // 这个值可以根据需要进行调整
	    // 分配推理请求向量的大小
	    infer_requests.resize(num_requests); // 调整向量大小
	    // 创建推理请求
	    for (size_t i = 0; i < num_requests; ++i) {
	        // 创建推理请求并存储在向量的相应位置
	        infer_requests[i] = compiled_model.create_infer_request();
	    }
	    // 使用空数据进行推理测试，如果测试失败则返回 false
	    if (!test_inference_with_empty_data(
	        infer_requests[0],
	        compiled_model.input().get_shape(),
	        compiled_model.input().get_element_type())) {
	        std::cerr << "Model inference test failed with empty data." << std::endl;
	        return false;
	    }
	    return true;
	}

	/**
	 * @brief 获取模型的批量大小。
	 *
	 * @return 返回模型信息中存储的批量大小。
	 */
	int Yolo_VINO::BatchSize() const {
		return o_m_info.batch_size;
	}


	/**
	 * @brief 执行目标检测。
	 * 使用VINO进行图像的预处理、推理和后处理，以获取最终的检测结果。
	 *
	 * @param p_imgs 指向输入图像数据的指针。
	 * @param num 输入图像的数量。
	 * @param vv_results 用于存储检测结果的向量的引用。
	 */
	void Yolo_VINO::Detect(
	    st_dlm_data* p_imgs,
	    int num,
	    std::vector<std::vector<DetRes_VINO>>& vv_results)
	{
	    // 创建一个向量来存储临时的cv::Mat对象
	    std::vector<cv::Mat> mats(num);

	    // 遍历p_imgs数组，将每个st_dlm_data中的图像数据转换为cv::Mat
	    for (int i = 0; i < num; ++i) {
	        const st_dlm_data& img_data = p_imgs[i];
	        const cv::Mat mat = img_data.img.rgb8; // 假设img.rgb8已经是cv::Mat类型

	        mats[i] = mat; // 将转换后的Mat存储到向量中
	        std::cout << "Image " << i << ": " << mats[i].rows << "x" << mats[i].cols << std::endl;
	    }

	    // 获取网络的宽度、高度和通道数。
	    int net_w = o_m_info.image_width,
	        net_h = o_m_info.image_height,
	        net_c = o_m_info.input_channel;

	    // 根据网络的通道数设置Scalar对象
	    cv::Scalar mean_val(0, 0, 0); // 默认均值为0
	    if (net_c == 3) {
	        // 如果网络需要3个通道，设置BGR均值
	        mean_val = cv::Scalar(0, 0, 0);
	    } else {
	        // 如果网络只需要一个通道，设置单一均值
	        mean_val = cv::Scalar(0);
	    }


	    // 根据网络的尺寸和均值预处理图像
	    std::vector<cv::Mat> prepared_mats = preprocess_image(
	        mats,
	        net_w,
	        net_h,
	        mean_val);

	    std::cout << "prepared_mats size " << prepared_mats.size() << std::endl;

	    std::vector<ov::Tensor> V_Tensor; // 用于存储输出张量
	    vv_results.clear(); // 清空检测结果向量

	    for (size_t i = 0; i < prepared_mats.size(); ++i) {
	        // 获取模型输入端口
	        Output<const ov::Node> input_port = compiled_model.input();
	        // 创建输入张量
	        ov::Tensor input_tensor(
	            input_port.get_element_type(),
	            input_port.get_shape(),
	            prepared_mats[i].ptr(0)); // 创建输入张量

	        // 设置输入张量
	        infer_requests[0].set_input_tensor(input_tensor);

	        // 执行推理
	        infer_requests[0].infer();

	        // 获取推理结果
	        ov::Tensor output = infer_requests[0].get_output_tensor(0);

	        // 获取输出张量的形状并打印
	        std::cout << "Output tensor shape for image " << i << ": ";
	        for (size_t j = 0; j < output.get_shape().size(); ++j) {
	            std::cout << output.get_shape()[j] << (j < output.get_shape().size() - 1 ? " x " : "");
	        }
	        std::cout << std::endl;

	        // 将输出张量存储起来
	        V_Tensor.push_back(output);

	        // 对推理结果进行后处理，获取检测结果
	        std::vector<DetRes_VINO> detection_results = postProcess_yolov6(
	            mats[i],
	            output,
	            o_m_info.obj_threshold,
	            o_m_info.nms_threshold);
	        // 将检测结果添加到vv_results中
	        vv_results.push_back(detection_results);
	    }

	    // 遍历每个图像的检测结果，并进行可视化保存
	    for (size_t i = 0; i < vv_results.size(); ++i) {
	        const auto& detection_results = vv_results[i];
	        const cv::Mat& original_img = mats[i]; // 获取原始图像

	        // 调用函数，绘制并保存检测结果
	        std::string output_file_path = "output_image_" + std::to_string(i) + ".jpg";
	        visualize_and_save_yolov6(original_img, detection_results, output_file_path);
	    }
	}

	/**
	 * @brief 对一批图像的推理结果进行后处理。
	 *
	 * 该函数接受一批图像和对应的模型输出张量，然后对每对图像和输出张量进行处理。
	 * 对于每张图像，它调用 `postProcess_yolov6` 函数来解码检测框、应用非极大值抑制（NMS），
	 * 并将结果存储在 `vv_results` 向量中。
	 *
	 * @param img_batch 一批图像的向量。
	 * @param data_tensors 与图像对应的模型输出张量的向量。
	 * @param prob_threshold_ 概率阈值，用于过滤低置信度的检测框。
	 * @param nms_threshold_ 非极大值抑制的阈值，用于消除重叠的检测框。
	 * @return 返回一个二维向量，包含所有图像的检测结果。
	 */
	std::vector<std::vector<Yolo_VINO::DetRes_VINO>> Yolo_VINO::postProcess_batch(
		const std::vector<cv::Mat> &img_batch, // 输入的图像批次
		const std::vector<ov::Tensor> &data_tensors,
		const float prob_threshold_,
		const float nms_threshold_)
	{
		std::vector<std::vector<DetRes_VINO>> batch_results; // 存储批次结果的向量

		// 确保输入的图像和张量数量相同
		assert(img_batch.size() == data_tensors.size());

		// 遍历图像和张量，对每对进行处理
		for (size_t i = 0; i < img_batch.size(); ++i) {
			const cv::Mat& src_img = img_batch[i];
			const ov::Tensor& data_tensor = data_tensors[i];

			// 调用单个图像的后处理函数
			std::vector<DetRes_VINO> results =
				postProcess_yolov6(
					src_img,
					data_tensor,
					prob_threshold_,
					nms_threshold_);

			// 将结果添加到批次结果中
			batch_results.push_back(results);
		}

		return batch_results; // 返回批次的检测结果
	}

	/**
	 * @brief 对单张图像的模型输出进行后处理，包括解码检测框和应用非极大值抑制（NMS）
	 * @param src_img 输入的单张图像
	 * @param output 模型输出的数据指针
	 * @return std::vector<DetectRes> 包含所有检测结果的向量
	 *
	 * 这个函数对单张图像的模型输出进行解码，将输出数据转换为检测框，并应用非极大值抑制来去除重叠的检测框。
	 */
	std::vector<Yolo_VINO::DetRes_VINO> Yolo_VINO::postProcess_yolov6(
	    const cv::Mat &src_img, // 输入的单张图像
	    const ov::Tensor& data_tensor,
	    const float prob_threshold_,
	    const float nms_threshold_
	    )  { // 模型输出的数据指针

	    auto* output = data_tensor.data<float>(); // 获取输出数据

	    std::vector<DetRes_VINO> result; // 存储检测结果的向量，确保类型与返回值类型一致
	    // 计算图像尺寸与模型输入尺寸的比例
	    const float ratio =
    		std::max(
    			static_cast<float>(src_img.cols) / static_cast<float>(o_m_info.image_width),
    			static_cast<float>(src_img.rows) / static_cast<float>(o_m_info.image_height));
	    // 计算垂直方向和水平方向的填充
	    const float scale = std::min(
    		static_cast<float>(o_m_info.image_width) / static_cast<float>(src_img.cols),
    		static_cast<float>(o_m_info.image_height) / static_cast<float>(src_img.rows));

	    float padd_h = static_cast<float>(
    		std::round((static_cast<float>(o_m_info.image_height) -
    			static_cast<float>(src_img.rows) * scale) / 2));
	    float padd_w = static_cast<float>(
    		std::round((static_cast<float>(o_m_info.image_width) -
    			static_cast<float>(src_img.cols) * scale) / 2));

	    for (auto& grid : m_grids) { // 遍历每个网格
	        for (int h = 0; h < grid[0]; h++) { // 遍历网格的高度
	            for (int w = 0; w < grid[1]; w++) { // 遍历网格的宽度
	                float* row = output + ((h * grid[1] + w) * (CATEGORY + 5)); // 获取当前网格的输出数据行
	                DetRes_VINO box{}; // 创建一个检测框结构体，确保类型与返回值类型一致
	                // 找到最大概率的类别
	                auto max_pos = std::max_element(row + 5, row + CATEGORY + 5);
	                box.prob = row[4] * row[max_pos - row]; // 计算检测框的概率
	                if (box.prob < prob_threshold_) // 如果概率低于阈值，则跳过
	                    continue;
	                box.classes = static_cast<int>(max_pos - row - 5); // 存储类别索引
	                // 解码检测框坐标
	                box.x = (row[0] - padd_w) * ratio; // x坐标
	                box.y = (row[1] - padd_h) * ratio; // y坐标
	                box.w = row[2] * ratio; // 宽度
	                box.h = row[3] * ratio; // 高度
	                result.push_back(box); // 将检测框添加到结果向量
	            }
	        }
	    }
	    cout<<"result before NMS:"<<result.size()<<endl;
	    // 应用非极大值抑制
	    NmsDetect(
    		result,
    		nms_threshold_); // 对检测结果应用NMS
	    cout<<"result after NMS:"<<result.size()<<endl;

	    return result; // 返回检测结果
	}

	/**
	* @brief 绘制并保存 YOLOv6 的检测结果。
	*
	* 该函数将 YOLOv6 检测结果在原始图像上进行可视化，包括绘制边界框和标签，
	* 并将最终结果保存到指定的输出路径。
	*
	* @param original_img 原始图像，用于绘制检测结果。
	* @param detection_results 检测结果的向量，包括检测框、类别 ID、概率等信息。
	* @param output_file_path 输出文件路径，用于保存可视化结果的图像。
	*/
	void visualize_and_save_yolov6(
	    const cv::Mat& original_img,
	    const std::vector<Yolo_VINO::DetRes_VINO>& detection_results,
	    const std::string& output_file_path)
	{
	std::vector<cv::Scalar> class_colors; // 类别颜色

	class_colors.resize(1); // 根据类别数量初始化颜色数组
	srand(static_cast<int>(time(nullptr))); // 设置随机种子
	for (cv::Scalar& class_color : class_colors) {
		class_color = cv::Scalar(rand() % 255, rand() % 255, rand() % 255); // 为每个类别生成随机颜色
	}
	cout << "class_colors num:" <<class_colors.size() << endl;
    // 如果原始图像为空，输出错误信息并返回
    if (original_img.empty()) {
        std::cerr << "Error: Original image is empty." << std::endl;
        return;
    }

    // 克隆原始图像，避免修改原图
    cv::Mat img = original_img.clone();

    // 遍历检测结果的向量
    for (const auto&[classes, x, y, w, h, prob] : detection_results) {
        // 根据检测框的中心点和宽高计算边界框的矩形区域
        cv::Rect box(x - w / 2, y - h / 2, w, h);

        // 检查类别 ID 是否在类别名称范围内
        if (const int class_id = classes; class_id >= 0 && class_id < static_cast<int>(1)) {
            // 绘制边界框，颜色为对应类别的颜色
            cv::rectangle(img, box, class_colors[class_id % class_colors.size()], 2, 8);

            // 创建标签字符串，包括类别名称和检测概率（保留小数点后4位）
            std::string label = "1:" + std::to_string(prob).substr(0, 4);

            // 获取标签文本的尺寸
            const cv::Size textSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, 0);

            // 创建文本框，用于显示类别标签
            const cv::Rect textBox(box.tl().x, box.tl().y - 15, textSize.width, textSize.height + 5);

            // 绘制文本框，使用类别对应的颜色填充
            cv::rectangle(img, textBox, class_colors[class_id % class_colors.size()], cv::FILLED);

            // 在文本框中绘制类别标签，字体颜色为白色
            cv::putText(
                img, label,
                cv::Point(box.tl().x, box.tl().y - 5),
                cv::FONT_HERSHEY_SIMPLEX, 0.5,
                cv::Scalar(255, 255, 255));
        } else {
            // 如果类别 ID 超出范围，输出错误信息
            std::cerr << "Error: Class ID out of range." << std::endl;
        }
    }

    // 尝试将可视化后的图像保存到指定路径
    if (!cv::imwrite(output_file_path, img)) {
        std::cerr << "Error: Failed to save image to " << output_file_path << std::endl; // 如果保存失败，输出错误信息
    } else {
        std::cout << "Detection results have been saved to: " << output_file_path << std::endl; // 保存成功时输出保存路径
    }
}


}//namespace yolo_v6